{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Data Classification with NeMo Curator's `DomainClassifier`\n",
    "\n",
    "This notebook demonstrates the use of NeMo Curator's `DomainClassifier`. The [domain classifier](https://huggingface.co/nvidia/domain-classifier) is used to classify the domain of a text. It helps with data annotation, which is useful in data blending for foundation model training. Please refer to the NemoCurator Domain Classifier Hugging Face page for more information about the domain classifier, including its output labels, here: https://huggingface.co/nvidia/domain-classifier.\n",
    "\n",
    "This tutorial requires at least 1 NVIDIA GPU with:\n",
    "  - Volta™ or higher (compute capability 7.0+)\n",
    "  - CUDA 12.x\n",
    "\n",
    "Before running this notebook, see this [Installation Guide](https://docs.nvidia.com/nemo/curator/latest/admin/installation.html#admin-installation) page for instructions on how to install NeMo Curator. Be sure to use an installation method which includes GPU dependencies.\n",
    "\n",
    "For more information about the classifiers, refer to our [Distributed Data Classification](https://docs.nvidia.com/nemo/curator/latest/curate-text/process-data/quality-assessment/distributed-classifier.html) documentation page."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Silence Curator logs via Loguru\n",
    "import os\n",
    "\n",
    "os.environ[\"LOGURU_LEVEL\"] = \"ERROR\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following imports are required for this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from nemo_curator.core.client import RayClient\n",
    "from nemo_curator.pipeline import Pipeline\n",
    "from nemo_curator.stages.text.classifiers import DomainClassifier\n",
    "from nemo_curator.stages.text.io.reader.jsonl import JsonlReader\n",
    "from nemo_curator.stages.text.io.writer.jsonl import JsonlWriter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run a pipeline in NeMo Curator, we must start a Ray cluster. This can be done manually (see the [Ray documentation](https://docs.ray.io/en/latest/ray-core/starting-ray.html)) or with Curator's `RayClient`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    ray_client = RayClient()\n",
    "    ray_client.start()\n",
    "except Exception as e:\n",
    "    msg = f\"Error initializing Ray client: {e}\"\n",
    "    raise RuntimeError(msg) from e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Read, Classification, and Write Stages\n",
    "\n",
    "Functions in NeMo Curator are called stages. For this tutorial, we will initialize 3 stages: a JSONL file reader, a domain classification stage, and a JSONL file writer.\n",
    "\n",
    "For this tutorial, we will create a sample JSONL file to use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_file_path = \"./input_data_dir\"\n",
    "\n",
    "# Create sample dataset for the tutorial\n",
    "text = [\n",
    "    \"Quantum computing is set to revolutionize the field of cryptography.\",\n",
    "    \"Investing in index funds is a popular strategy for long-term financial growth.\",\n",
    "    \"Recent advancements in gene therapy offer new hope for treating genetic disorders.\",\n",
    "    \"Online learning platforms have transformed the way students access educational resources.\",\n",
    "    \"Traveling to Europe during the off-season can be a more budget-friendly option.\",\n",
    "    \"Training regimens for athletes have become more sophisticated with the use of data analytics.\",\n",
    "    \"Streaming services are changing the way people consume television and film content.\",\n",
    "    \"Vegan recipes have gained popularity as more people adopt plant-based diets.\",\n",
    "    \"Climate change research is critical for developing sustainable environmental policies.\",\n",
    "    \"Telemedicine has become increasingly popular due to its convenience and accessibility.\",\n",
    "]\n",
    "df = pd.DataFrame({\"text\": text})\n",
    "\n",
    "try:\n",
    "    os.makedirs(input_file_path, exist_ok=True)\n",
    "    df.to_json(input_file_path + \"/data.jsonl\", orient=\"records\", lines=True)\n",
    "except Exception as e:\n",
    "    msg = f\"Error creating input file: {e}\"\n",
    "    raise RuntimeError(msg) from e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define the reader stage with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read existing directory of JSONL files\n",
    "read_stage = JsonlReader(input_file_path, files_per_partition=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The classifier stage is broken down under the hood into a tokenizer stage and a model inference stage. Tokenization is run on the CPU while model inference is run on the GPU. This means that behind the scenes, the `DomainClassifier` stage is actually being broken down into 2 stages (some parameters and details omitted to avoid complexity, please refer to the documentation for more details):\n",
    "\n",
    "```python\n",
    "class TokenizerStage:\n",
    "    self.resources = Resources(cpus=1)\n",
    "    self.model_identifier = \"nvidia/domain-classifier\"\n",
    "    self.text_field = \"text\"\n",
    "    self.padding_side = \"right\"\n",
    "    ...\n",
    "class ModelStage:\n",
    "    self.resources = Resources(cpus=1, gpus=1)\n",
    "    self.model_identifier = \"nvidia/domain-classifier\"\n",
    "    self.model_inference_batch_size = 256\n",
    "    ...\n",
    "```\n",
    "\n",
    "Optionally, the classifier predictions may be filtered to include only texts with values listed in `filter_by`. If the `filter_by` parameter is set, then a third stage is added:\n",
    "\n",
    "```python\n",
    "def filter_by_category(self, value: str) -> bool:\n",
    "    return value in self.filter_by\n",
    "\n",
    "...\n",
    "\n",
    "if self.filter_by is not None and len(self.filter_by) > 0:\n",
    "    self.stages.append(Filter(filter_fn=self.filter_by_category, filter_field=...))\n",
    "```\n",
    "\n",
    "See the [API Reference](https://docs.nvidia.com/nemo/curator/latest/apidocs/stages/stages.text.classifiers.domain.html#stages.text.classifiers.domain.DomainClassifier) for more information about the `DomainClassifier` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the domain classifier\n",
    "classifier_stage = DomainClassifier()\n",
    "\n",
    "# If desired, you may filter your dataset with:\n",
    "# classifier_stage = DomainClassifier(filter_by=[\"Computers_and_Electronics\", \"Health\"])  # noqa: ERA001\n",
    "# See full list of domains here: https://huggingface.co/nvidia/domain-classifier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can define a stage for writing the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write results to a directory\n",
    "output_file_path = \"./domain_classifier_results\"\n",
    "\n",
    "# Use mode=\"overwrite\" to overwrite the output directory if it already exists\n",
    "# This helps to ensure that the correct output is written\n",
    "write_stage = JsonlWriter(output_file_path, mode=\"overwrite\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Pipeline\n",
    "\n",
    "In NeMo Curator, we use pipelines to run distributed data workflows using Ray. Pipelines take care of resource allocation and autoscaling to achieve enhanced performance and minimize GPU idleness.\n",
    "\n",
    "For the distributed data classifiers, we are able to achieve speedups by ensuring that model inference is run in parallel across all available GPUs, while other stages such as I/O, tokenization, and filtering are run across all available CPUs. This is possible because Curator pipelines are composable, which allows each stage in a pipeline to run independently and with its own specified hardware resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Pipeline(name='classifier_pipeline', stages=[jsonl_reader(JsonlReader), domain_classifier_classifier(DomainClassifier), jsonl_writer(JsonlWriter)])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier_pipeline = Pipeline(name=\"classifier_pipeline\", description=\"Run a classifier pipeline\")\n",
    "\n",
    "# Add stages to the pipeline\n",
    "classifier_pipeline.add_stage(read_stage)\n",
    "classifier_pipeline.add_stage(classifier_stage)\n",
    "classifier_pipeline.add_stage(write_stage)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Composability is also what allows a classifier to sit between pre-processing and post-processing stages. Typical text pre-processing add-ons include text normalization (lowercasing, URL/email removal, Unicode cleanup) and language identification and filtering (to keep only target languages). A full pipeline may look something like:\n",
    "\n",
    "```python\n",
    "pipeline = Pipeline(name=\"full_pipeline\")\n",
    "pipeline.add_stage(read_stage)                # reader (JSONL/S3/etc.)\n",
    "pipeline.add_stage(lang_id_stage)             # optional: language filter\n",
    "pipeline.add_stage(classifier_stage)          # classifier\n",
    "pipeline.add_stage(write_stage)               # writer (JSONL/Parquet)\n",
    "```\n",
    "\n",
    "# Run the  Classifier\n",
    "\n",
    "Let's run the full classifier pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the pipeline\n",
    "result = classifier_pipeline.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the pipeline ran to completion and the result was written to a JSONL file, we can shut down the Ray cluster with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    ray_client.stop()\n",
    "except Exception as e:  # noqa: BLE001\n",
    "    print(f\"Error stopping Ray client: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspect the Output\n",
    "\n",
    "The write stage returns a list of written files. We can read the output file as a Pandas DataFrame for inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>domain_pred</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Quantum computing is set to revolutionize the ...</td>\n",
       "      <td>Computers_and_Electronics</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Investing in index funds is a popular strategy...</td>\n",
       "      <td>Finance</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Recent advancements in gene therapy offer new ...</td>\n",
       "      <td>Health</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Online learning platforms have transformed the...</td>\n",
       "      <td>Jobs_and_Education</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Traveling to Europe during the off-season can ...</td>\n",
       "      <td>Travel_and_Transportation</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                text  \\\n",
       "0  Quantum computing is set to revolutionize the ...   \n",
       "1  Investing in index funds is a popular strategy...   \n",
       "2  Recent advancements in gene therapy offer new ...   \n",
       "3  Online learning platforms have transformed the...   \n",
       "4  Traveling to Europe during the off-season can ...   \n",
       "\n",
       "                 domain_pred  \n",
       "0  Computers_and_Electronics  \n",
       "1                    Finance  \n",
       "2                     Health  \n",
       "3         Jobs_and_Education  \n",
       "4  Travel_and_Transportation  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# For simplicity, we take the first written file from the writer stage\n",
    "# In real pipelines, the writer may return multiple files (shards) or objects\n",
    "result_file = result[0].data[0]\n",
    "\n",
    "result_df = pd.read_json(result_file, lines=True)\n",
    "result_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the predictions were generated as expected."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
